import torch
import torch.nn as nn
import torch.nn.functional as F

from detection.layers import weight_init


# mmdet style
class FPN(nn.Module):
    def __init__(self, cfg, in_channels, out_channels=256):
        super().__init__()
        self.in_channels = in_channels
        self.out_channels = out_channels
        self.num_ins = len(in_channels)
        self.num_outs = self.num_ins + 1    # add stride=64

        self.lateral_convs = nn.ModuleList()
        self.fpn_convs = nn.ModuleList()

        for i in range(self.num_ins):
            l_conv = nn.Conv2d(in_channels[i], out_channels, kernel_size=1, stride=1)
            fpn_conv = nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=1, padding=1)
            # init weight
            weight_init.c2_xavier_fill(l_conv)
            weight_init.c2_xavier_fill(fpn_conv)
            # add
            self.lateral_convs.append(l_conv)
            self.fpn_convs.append(fpn_conv)
        self._fuse_type = 'sum'
            
    def forward(self, inputs):
        assert len(inputs) == self.num_ins

        # build laterals
        laterals = [
            lateral_conv(inputs[i])
            for i, lateral_conv in enumerate(self.lateral_convs)
        ]

        # build top-down path
        num_laterals = len(laterals)
        for i in range(num_laterals - 1, 0, -1):
            prev_shape = laterals[i - 1].shape[2:]
            laterals[i - 1] += F.interpolate(laterals[i], size=prev_shape, mode='nearest')
        
        # build fpn(outputs)
        outs = [
            self.fpn_convs[i](laterals[i]) for i in range(num_laterals)
        ]

        # extra levels  (stride=64)
        outs.append(
            F.max_pool2d(outs[-1], kernel_size=1, stride=2)
        )

        return outs


# detectron2 style
# class FPN(nn.Module):
#     def __init__(self, cfg, in_channels, out_channels=256):
#         super().__init__()
#         self.in_channels = in_channels
#         self.out_channels = out_channels
#         self.num_ins = len(in_channels)
#         self.num_outs = self.num_ins + 1    # add stride=64

#         lateral_convs = []
#         fpn_convs = []

#         for i in range(self.num_ins):
#             l_conv = nn.Conv2d(in_channels[i], out_channels, kernel_size=1, stride=1)
#             fpn_conv = nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=1, padding=1)
#             # init weight
#             weight_init.c2_xavier_fill(l_conv)
#             weight_init.c2_xavier_fill(fpn_conv)
#             # add
#             lateral_convs.append(l_conv)
#             fpn_convs.append(fpn_conv)
#         self.lateral_convs = nn.ModuleList(lateral_convs[::-1])
#         self.fpn_convs = nn.ModuleList(fpn_convs[::-1])
#         self._fuse_type = 'sum'

            
#     def forward(self, inputs):
#         assert len(inputs) == self.num_ins

#         results = []
#         prev_features = self.lateral_convs[0](inputs[-1])
#         results.append(self.fpn_convs[0](prev_features))

#         for idx, (lateral_conv, output_conv) in enumerate(
#             zip(self.lateral_convs, self.fpn_convs)
#         ):
#             if idx > 0:
#                 features = inputs[-idx - 1]
#                 top_down_features = F.interpolate(prev_features, scale_factor=2.0, mode="nearest")
#                 lateral_features = lateral_conv(features)
#                 prev_features = lateral_features + top_down_features
#                 if self._fuse_type == "avg":
#                     prev_features /= 2
#                 results.insert(0, output_conv(prev_features))
        
#         maxpooled_in_feature = results[-1]
#         results.append(F.max_pool2d(maxpooled_in_feature, kernel_size=1, stride=2, padding=0))

#         return results